import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from sklearn.metrics import accuracy_score
from PIL import Image
import matplotlib.pyplot as plt

# 示例数据：生成一些简单的训练和测试数据（例如正弦曲线）
n = 1000
x_train = torch.unsqueeze(torch.linspace(-5, 5, n), dim=1)
y_train = torch.sin(x_train) + 0.2 * torch.randn(n, 1)

x_test = torch.unsqueeze(torch.linspace(-5, 5, 300), dim=1)
y_test = torch.sin(x_test) + 0.2 * torch.randn(300, 1)


# 定义一个简单的神经网络模型
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 8)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(8, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x


# 初始化模型和训练参数
model = SimpleNet()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.02)

# 初始化TensorBoard的SummaryWriter
writer = SummaryWriter(log_dir='runs/example_torch_tensorboard')

# 训练循环
num_epochs = 1000
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(x_train)
    loss = criterion(outputs, y_train)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # TensorBoard记录各种数据
    writer.add_scalar('training_loss', loss.item(), epoch)
    if epoch % 10 == 0:  # 每10个epoch添加其他指标
        with torch.no_grad():
            # 验证 loss
            val_outputs = model(x_test)
            val_loss = criterion(val_outputs, y_test)
            writer.add_scalar('validation_loss', val_loss.item(), epoch)

            # 图表和权重直方图
            writer.add_histogram('model_weights', model.state_dict()['fc1.weight'], epoch)
            writer.add_histogram('model_grads', model.fc1.weight.grad, epoch)

            # 添加模型结构到TensorBoard
            dummy_input = torch.randn(10, 1)
            output = model(dummy_input)
            writer.add_graph(model, (dummy_input,), epoch)

            # 添加文本注释
            writer.add_text('hyperparameters',
                            f'learning_rate={optimizer.param_groups[0]["lr"]}\nbatch_size=32',
                            epoch)

            # 添加图像示例（假设y_pred和y真实值）
            plt.figure(figsize=(10, 5))
            plt.plot(x_train.numpy(), y_train.numpy(), 'r-', label='True')
            plt.plot(x_train.numpy(), outputs.detach().numpy(), 'b-', label='Predicted')
            writer.add_figure('train_prediction_plots', plt.gcf(), epoch)

        # 添加模型权重和梯度的标量指标（例如最大值、最小值等）
        weights = model.fc1.weight.data.numpy()
        writer.add_scalar('fc1_weights_max', np.max(weights), epoch)
        writer.add_scalar('fc1_weights_min', np.min(weights), epoch)

# 最后关闭SummaryWriter
writer.close()
